Blackwell Programming for the Masses With OpenAI Triton

Phil Tillet | OpenAI
dePaul Miller | NVIDIA

目录

演讲大纲 (Talk Outline)

Triton 编程语言 (The Triton Programing Language)

动机：CUDA 的权衡 (Motivations: CUDA trade-offs)

动机：XLA/Torch 的权衡 (Motivations: XLA/Torch trade-offs)

动机：Triton 的权衡 (Motivations: Triton trade-offs)

生产力 / 性能帕累托前沿 (Productivity / Performance Pareto Frontier)

Page 6: 展示了不同编程环境在生产力和性能之间的权衡关系图。纵轴为生产力，横轴为性能。从左上到右下，依次为Torch、Triton、CUDA、PTX，表明生产力逐渐降低而性能逐渐提高。
Page 6: 展示了不同编程环境在生产力和性能之间的权衡关系图。纵轴为生产力，横轴为性能。从左上到右下，依次为Torch、Triton、CUDA、PTX，表明生产力逐渐降低而性能逐渐提高。

Triton 编程模型 (The Triton Programming Model)

概述 (Overview)

矩阵乘法 (Matrix Multiplication)

概述 (Overview)

以下是使用 Triton 实现矩阵乘法的代码示例。
Page 8: Triton JIT 编译的矩阵乘法 Python 代码。函数定义为 matmul，包含 prologue（序言）、main loop（主循环）和 epilogue（结尾）部分。

性能 (Performance) (FP16; M=N=8192; cuBLAS 12.8)

下图比较了 Triton 和 cuBLAS 在矩阵乘法上的性能。
Page 9: Triton 与 cuBLAS 在 FP16 矩阵乘法上的性能对比图。Triton 的性能曲线（蓝色）始终低于 cuBLAS（绿色），尤其是在 K 值较小时。

软件流水线 (Software Pipelining)

持久化矩阵乘法 (Persistent Matrix Multiplication)

概述 (Overview)

以下是持久化矩阵乘法的 Triton 实现代码。
Page 11: 持久化矩阵乘法的 Triton JIT 代码。与基础版本相比，该代码通过循环处理多个 tile_id 来实现持久化。

软件流水线 (Hopper)

H100 性能 (FP16; M=N=8192; cuBLAS 12.8)

下图展示了在 H100 上的性能表现。
Page 13: H100 上持久化矩阵乘法的性能对比图。Triton 的性能（蓝色）与 cuBLAS（绿色）非常接近，但在 K 值较小时仍然较慢，如图中红色框所示。

软件流水线 (Blackwell)

下图对比了 Hopper 和 Blackwell 架构上的软件流水线。
Page 15: Hopper 与 Blackwell 软件流水线对比图。上图为 Hopper，存在 Epilogue 导致的阻塞。下图为 Blackwell，通过新的硬件特性，流水线更加紧凑，几乎消除了气泡，显著提高了效率。
让我们尝试在 Hopper 上运行 matmul 内核，然后……看看会发生什么？

GB200 性能表现

下图展示了在 GB200 平台上，使用 FP16 精度，当 M=N=8192 时，Triton 和 cuBLAS 12.8 在矩阵乘法上的性能对比。图中显示，Triton 的性能（蓝色曲线）显著低于 cuBLAS（绿色曲线）。cuBLAS 的性能大约在 1400 TFLOPS，而 Triton 的性能大约在 900 TFLOPS。
Page 16

瓶颈分析

从下方的流水线图中可以看出，在加载 Load A₂ 和 Load B₂ 的同时，张量核心（Tensor Cores）正在执行 MMA₀。但在 MMA₀ 完成后，由于 Load A₂ 和 Load B₂ 尚未完成，张量核心进入空闲状态，直到数据准备好才能开始执行 MMA₁。这种空闲状态表明计算单元未被充分利用。
Page 17
- 解决方案：尝试将流水线深度增加到 num_stages=4，观察其效果。

增加流水线深度

下图对比了 num_stages=3 和 num_stages=4 的情况。
- 当 num_stages=4 时，流水线图看起来更好。通过增加流水线深度，可以更好地隐藏数据加载的延迟，使得 MMA 操作能够连续执行，从而减少或消除了张量核心的空闲时间。
Page 18
- 让我们运行一下，看看实际结果如何。

Blackwell 平台性能问题... 哦不！

当尝试在 Blackwell 平台上运行时，程序抛出了资源不足的错误。
Page 19
错误信息如下：
triton.runtime.errors.OutOfResources: out of resource: shared memory. Required: 262176, Hardware limit: 232448. Reducing block sizes or 'num_stages' may help.
这表明增加流水线深度（num_stages）导致共享内存（shared memory）的需求超过了硬件限制。

结尾部分拆分 (Epilogue Splitting)

为了解决共享内存不足的问题，同时保持较高的流水线深度，可以采用“结尾部分拆分”（Epilogue Splitting）的策略。以下是实现该策略的 Triton JIT 代码示例。关键在于 tl.split(acc)，它将累加器 acc 拆分，允许以更小的块进行存储，从而减少对资源的瞬时需求。
Page 20

应用优化后的 GB200 性能

在应用了结尾部分拆分等优化后，Triton 的性能得到了显著提升。从下图中可以看到，Triton 的性能曲线（蓝色）现在更接近 cuBLAS（绿色），尤其是在 K 值较大时。尽管仍略低于 cuBLAS，但差距已经大幅缩小。
Page 21

注意事项与未来工作 (Caveats)

Blackwell

概述 (Overview)

再次探讨 Blackwell 上的 Group GEMM

寻找唾手可得的优化点 (Finding Low-Hanging Fruits)

修复网格大小 (Fixing the Grid Size)

@triton.autotune(configs=[
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'NUM_SM': 128}))
])
def num_sms():
    return torch.cuda.get_device_properties('cuda').multi_processor_count

@triton.autotune(configs=[
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'NUM_SM': num_sms()}),
])

与持久化 GEMM (Persistent GEMM) 的比较

调试性能差距：Nsight Systems

调试性能差距：Nsight Compute

与基准进行比较！
- 通过在 Nsight Compute 中运行两个内核来进一步挖掘。
- 我们可以右键点击 Persistent GEMM，将其设置为比较基准。
- 这证实了 Persistent GEMM 更好地利用了张量核心 (Tensor Cores, TC)。
- 我们需要更好地利用第 5 代张量核心。
Page 36 中的 Nsight Compute 管道利用率图表。蓝色条代表我们的 Group GEMM 内核，绿色条代表 Persistent GEMM。绿色条在 TC 和 Tensor (All) 上的利用率远高于蓝色条。

NVIDIA Blackwell 架构

Blackwell 性能 101

triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'NUM_SM': num_sms()}),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'NUM_SM': num_sms()}),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'NUM_SM': num_sms()}),
Page 38 中的示意图，展示了 A (128x64)、B (64x256)、C (128x256) 的 tile 尺寸。
Page 38 中的示意图，展示了 A (128x64)、B (64x256)、C (128x256) 的 tile 尺寸。

性能结果

进一步优化

Triton 中的 TMA

a_desc = tl.make_tensor_descriptor(
    a_ptr,
    shape=[group_shape_m, group_shape_k],
    strides=[group_stride_m, 1], # row major tensor
    block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K], # copy block size
)
#
a_desc.load([offs_m, offs_k]) # load with starting coordinate
Page 41 的示意图，展示了 TMA 如何通过描述符将全局内存张量中的一个块加载到共享内存中。
Page 41 的示意图，展示了 TMA 如何通过描述符将全局内存张量中的一个块加载到共享内存中。

性能增益

更进一步 (Going Further)

致谢